from services.AIServiceBase import AIServiceBase
from typing import Optional, Dict, Any
import json
from zhipuai import ZhipuAI
import asyncio


class ZhipuService(AIServiceBase):
    def __init__(self, api_key: str, config: Dict = None):
        super().__init__(api_key, config)
        self.client = ZhipuAI(
            api_key=api_key,
        )
        self.config = config if isinstance(config, dict) else {}

    async def generate_text(self, prompt: str) -> Optional[str]:
        """
        核心文本生成
        :param prompt: 提示词
        :return:
        """
        try:
            model_params = {
                "temperature": self.config.get("temperature", 0.7),
                "top_p": self.config.get("top_p", 0.9),
                "max_tokens": self.config.get("max_tokens", 1000),
            }
            response = await asyncio.to_thread(
                self.client.chat.completions.create,
                model="glm-4v-flash",
                messages=[
                    {"role": "system", "content": "You are a helpful assistant"},
                    {"role": "user", "content": prompt}
                ],
                **model_params
            )
            return response.choices[0].message.content
        except Exception as e:
            self.logger.error(f"Zhipu生成失败{str(e)}", exc_info=True)
            return None

    async def continue_writing(self, text: str) -> Optional[str]:
        self.logger.info("开始续写")
        prompt = f"""请续写以下内容
            {text}
            续写要求：
            1.保持原有的风格和语气
            2.续写200-300字
            3.确保内容连贯
            """
        result = await self.generate_text(prompt)
        return result

    async def ask_question(self, question: str) -> Optional[str]:
        """问答功能实现"""
        self.logger.info(f"开始回答问题{question[:100]}...")
        prompt = f"""请回答以下问题
            {question}
            回答要求：
            1.提供详细解释
            2.必要时给出示例
            3.保持专业性
            """
        self.logger.debug(f"问题提示词：{prompt}")
        result = await self.generate_text(prompt)
        self.logger.info("问题已回答")
        return result

    async def expand_text(self, text: str) -> Optional[str]:
        """文本扩展功能实现"""
        self.logger.info("开始扩写")
        prompt = f"""请扩写以下内容
        {text}
        扩写要求：
        1.保持核心内容不变
        2.请添加细节和解释
        3.使内容更加丰富
        """
        self.logger.debug(f"扩写提示词：{prompt}")
        result = await self.generate_text(prompt)
        self.logger.info("内容以扩写")
        return result

    async def analyze_code(self, code: str) -> Optional[Dict[str, Any]]:
        """代码分析功能实现"""
        self.logger.info("开始代码分析")
        prompt = f"""请分析一下代码并返回JSON格式的分析结果

        分析要求：
        1.代码质量评分（0-100）
        2.潜在问题列表
        3.改进的建议

        代码：{code}

        返回格式：
        {{
            "quality_score":0-100,
            "issues":["问题1","问题2"],
            "suggestions":["建议1","建议2"]
        }}
        """
        self.logger.debug(f"代码分析提示词：{prompt}")
        result = await self.generate_text(prompt)
        if result:
            try:
                return json.loads(result)
            except json.JSONDecodeError:
                self.logger.error("分析结果转换为json失败")
                return None
        self.logger.warning("分析结果为空")
        return None

    async def summarize_text(self, text: str, max_length: int = 150) -> Optional[str]:
        """文本摘要功能的实现"""
        prompt = f"""请对以下文本进行摘要，摘要长度不超过{max_length}字:{text},如果需要摘要的内容过短，不要强行续写"""
        return await self.generate_text(prompt)

    async def translate_text(self, text: str, target_language: str = "en") -> Optional[str]:
        """文本翻译功能实现"""
        prompt = f"""请将以下文本翻译成{target_language}:{text}"""
        return await self.generate_text(prompt)

    async def convert_style(self, text: str, target_style: str = "informal") -> Optional[str]:
        """文本风格转换功能实现"""
        prompt = f"""请将以下文本风格转换成{target_style}:{text},不要改变文本大意"""
        return await self.generate_text(prompt)